import pandas as pd
import numpy as np
from sklearn.linear_model import LinearRegression
import matplotlib.pyplot as plt
import warnings

warnings.filterwarnings("ignore")

# 读取数据
name = '附件2.xlsx'  # Excel文件名称，包含原始数据
data = pd.read_excel(name, header=None, names=['C1', 'C2', 'C3', 'C4'])  # 从Excel文件读取数据，数据包含C1、C2、C3和C4四列

# 筛选出出现至少两次的(C1, C3)组合
com = data.groupby(['C1', 'C3']).size().reset_index(name='count')  # 计算每个C1和C3组合的出现次数
com = com[com['count'] >= 2][['C1', 'C3']]  # 过滤出出现次数至少为2次的C1和C3组合

# 线性回归预测过程
def pre(data_group):  # 进行线性回归预测的函数
    x = data_group['time_index'].values.reshape(-1, 1)  # 获取时间索引作为自变量
    y = data_group['C4'].values  # 获取C4列作为因变量
    linreg = LinearRegression()  # 创建线性回归模型
    linreg.fit(x, y)  # 拟合模型
    next_time = x[-1][0] + 10  # 假设下一个时间点是当前时间点加10分钟
    return linreg.predict(np.array([[next_time]]))[0]  # 预测下一个时间点的C4值

# 处理预测数据
def pro_data(data, com):  # 处理数据并进行预测的函数
    results = []  # 存储结果的列表
    for _, row in com.iterrows():  # 遍历每个C1和C3的组合
        current_target, current_time = row['C1'], row['C3']  # 获取当前的C1和C3组合
        selected_data = data[(data['C1'] == current_target) & (data['C3'] == current_time)].sort_values(by='C2')  # 筛选出对应的数据并按C2排序

        # 时间格式转换和索引计算
        selected_data['C2'] = pd.to_datetime(selected_data['C2'], format='%Y/%m/%d %H:%M')  # 将C2列转换为日期时间格式
        selected_data['time_index'] = (selected_data['C2'] - selected_data['C2'].iloc[0]).dt.total_seconds() / 60  # 计算时间索引，单位为分钟

        # 预测处理
        try:
            lin_pred = pre(selected_data)  # 调用线性回归预测函数
        except Exception:
            lin_pred = np.nan  # 预测失败时返回NaN
            print(f"线性预测失败 {current_target} - 时间 {current_time}")  # 打印错误信息

        # 将每个组合的预测结果存储到结果列表中
        results.append({
            'target_str': current_target,  # 当前目标C1
            'selected_time': current_time,  # 当前时间C3
            'Linear_pred': lin_pred  # 预测的C4值
        })
    return pd.DataFrame(results)  # 返回结果数据框

# 预测数据处理
result = pro_data(data, com)  # 获取处理后的预测结果

# 添加日期列
def date(x):  # 根据预测值添加日期的函数
    return '2024/12/14' if isinstance(x, float) and x >= 0.85 else '2024/12/15'  # 依据预测值分配日期

result['date'] = result['selected_time'].apply(date)  # 根据预测值为每行添加日期

# 处理预测值
result['Linear_pred'] = result['Linear_pred'].round().fillna(0).astype(int)  # 对预测值进行四舍五入并填充NaN值为0，转换为整数
result['Linear_pred'] = result['Linear_pred'].apply(lambda x: max(x, 0))  # 保证预测值不小于0

# 保存10分钟粒度结果
result.to_excel('预测结果_10分钟.xlsx', index=False)  # 保存10分钟粒度的预测结果到Excel文件
print('已保存预测结果_10分钟.xlsx')  # 输出提示

# 汇总每日预测总量
sum = result.groupby('target_str')['Linear_pred'].sum().reset_index()  # 按target_str分组并对Linear_pred进行求和，得到每日总量
sum['date'] = '2024/12/15'  # 设置日期为2024/12/15
sum.rename(columns={'Linear_pred': 'Linear_pred_sum'}, inplace=True)  # 重命名预测列为Linear_pred_sum

# 保存汇总结果
sum.to_excel('预测结果_1天.xlsx', index=False)  # 保存1天的汇总结果到Excel文件
print('已保存预测结果_1天.xlsx')  # 输出提示

# 读取预测结果
result = pd.read_excel('预测结果_10分钟.xlsx')  # 读取10分钟粒度的预测结果
sum = pd.read_excel('预测结果_1天.xlsx')  # 读取1天汇总的预测结果

# 设置中文字体支持（根据需要可以注释）
plt.rcParams['font.sans-serif'] = ['SimHei']  # 配置中文字体
plt.rcParams['axes.unicode_minus'] = False  # 解决负号显示问题

# 图1：10分钟粒度散点图
plt.figure(figsize=(12, 6))  # 设置图表大小
for name, group in result.groupby('target_str'):  # 根据target_str分组绘制散点图
    plt.scatter(group['selected_time'], group['Linear_pred'], label=name)  # 绘制每个组的散点图
plt.title('10分钟粒度预测趋势散点图')  # 图表标题
plt.xlabel('时间点')  # x轴标签
plt.ylabel('预测值')  # y轴标签
plt.xticks(rotation=45)  # x轴标签旋转45度
plt.tight_layout()  # 自动调整布局以防止标签重叠
plt.savefig('图1_10分钟预测趋势.png')  # 保存图表
plt.show()  # 显示图表

# 图2：1天汇总预测折线图
plt.figure(figsize=(10, 5))  # 设置图表大小
plt.plot(sum['target_str'], sum['Linear_pred_sum'], marker='o')  # 绘制折线图
plt.title('1天预测总量趋势图')  # 图表标题
plt.xlabel('目标对象')  # x轴标签
plt.ylabel('总预测值')  # y轴标签

# 获取所有 x 轴标签
x_labels = sum['target_str'].tolist()
# 每隔一定间隔选取标签，这里设置间隔为 10，可根据实际情况调整
step = 10
selected_indices = range(0, len(x_labels), step)  # 选择要显示的x轴标签的索引
selected_labels = [x_labels[i] for i in selected_indices]  # 根据索引获取标签

# 设置 x 轴刻度
plt.xticks(selected_indices, selected_labels, rotation=45)  # 设置x轴刻度并旋转标签
plt.tight_layout()  # 自动调整布局
plt.savefig('图2_1天总量预测趋势.png')  # 保存图表
plt.show()  # 显示图表
